from sympy import Symbol, Eq, Ge, Abs, Function, Number, sin, cos
from modulus.pdes import PDES
from modulus.variables import Variables
import time
from modulus.solver import Solver
from modulus.dataset import TrainDomain, InferenceDomain
from modulus.data import Inference
from modulus.sympy_utils.geometry_2d import Rectangle, Line
from modulus.controller import ModulusController
import numpy as np
import math
import sys

def get_angle(theta, magnitude):
    return math.cos(theta)*magnitude, math.sin(theta)*magnitude

class Poisson_2D(PDES):
    name = 'Poisson_2D'
    def __init__(self):

        x, y = Symbol('x'), Symbol('y')

        alpha = Symbol('alpha')

        input_variables = {'x': x, 'y': y, 'alpha': alpha}

        phi = Function('phi')(*input_variables)
        
        self.equations = Variables()

        self.equations['u'] = phi.diff(x) 
        self.equations['v'] = phi.diff(y)
        self.equations['Poisson_2D'] = (phi.diff(x)).diff(x) + (phi.diff(y)).diff(y) 


obstacle_length = 0.10
height = 6*obstacle_length  
width = 6*obstacle_length

rec = Rectangle((-width / 2, -height / 2), (width / 2, height / 2))
obstacle = Line((0, 0), (0, obstacle_length), 1)
wake = Line((0, -3*obstacle_length), (0, 0), 1) 

obstacle.rotate(np.pi / 2)
wake.rotate(np.pi / 2)


geo = rec

x, y, alpha = Symbol('x'), Symbol('y'), Symbol('alpha')
param_ranges = {alpha, (-np.pi*10/180, np.pi*10/180)}
fixed_param_range = {alpha: lambda batch_size: np.full((batch_size, 1), np.random.uniform(-np.pi*10/180, np.pi*10/180))}

class PotentialTrain(TrainDomain):
    def __init__(self, **config):
        super(PotentialTrain, self).__init__()

#############################################################################################
        # I want to make the inlet velocity to be 10.0 m/s with an incidence angle of 4 degrees at the obstacle.
        # the inverse tan(v/u) gives me the required angle of incidence.
        # Drawing the scenario in comments below:
        #      +---------+
        #     /|/     \|/|
        #    //|// --- //|
        #   ///|/////////|
        #  ////+---------+
        #  //////////////
        #  / ////////////
        #    / //////////
        # where / is u + v such that tan-1(v/u) = x degrees(here i kept x as 4).
        u_x = 10*cos(alpha)
        u_y = 10*sin(alpha)
        flow_rate = u_x*width + u_y*height

        inletBC = geo.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y},
            batch_size_per_area=250,
            criteria=Eq(x, -width / 2),
            param_ranges ={**fixed_param_range},
            fixed_var=False
        )
        self.add(inletBC, name="Inlet")

        outletBC = geo.boundary_bc(
            outvar_sympy={"u":u_x , "v": u_y},
            batch_size_per_area=500,
            criteria=Ge(y/height+x/width, 1/2),
            param_ranges ={**fixed_param_range},
            fixed_var=False
        )
        self.add(outletBC, name="Outlet")

        bottomWall = geo.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y},
            batch_size_per_area=250,
            criteria=Eq(y, -height / 2),
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(bottomWall, name="BottomWall")

        obstacleLine = obstacle.boundary_bc(
            outvar_sympy={"u": u_x, "v": 0},
            batch_size_per_area=500,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100},
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(obstacleLine, name="obstacleLine")

   
        l = lambda x : (x)/(3*obstacle_length) 
        wakeLine = wake.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x)},
            batch_size_per_area=500,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, },
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(wakeLine, name="wakeLine")

        interior = geo.interior_bc(
            outvar_sympy={"Poisson_2D": 0},
            bounds={x: (-width / 2, width / 2), y: (-height / 2, height / 2)},
            lambda_sympy={
                "lambda_Poisson_2D": geo.sdf,
            },
            batch_size_per_area=2000,
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(interior, name="Interior")

        neighbourhood = geo.interior_bc(
            outvar_sympy={"Poisson_2D": 0},
            bounds={x: (-height / 3, height / 3), y: (-height / 8, height / 8)},
            lambda_sympy={
                "lambda_Poisson_2D": geo.sdf,
            },
            batch_size_per_area=2000,
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(neighbourhood, name="Neighbourhood")

class PotentialInference(InferenceDomain):
    def __init__(self, **config):
        super(PotentialInference, self).__init__()
        x, y, alpha = Symbol('x'), Symbol('y'), Symbol('alpha')
        interior = Inference(geo.sample_interior(10000, bounds={x: (-width / 2, width / 2), y: (-height / 2, height / 2)}, param_ranges={alpha: 1.0}), ['phi'])
        self.add(interior, name="Inference")

class PotentialSolver(Solver):
    train_domain = PotentialTrain
    inference_domain = PotentialInference

    def __init__(self, **config):
        super(PotentialSolver, self).__init__(**config)
        self.equations = (
            Poisson_2D().make_node()
        )
        flow_net = self.arch.make_node(
            name="flow_net", inputs=["x", "y", "alpha"], outputs=["phi"]
        )
        self.nets = [flow_net]

    @classmethod
    def update_defaults(cls, defaults):
        defaults.update(
            {
                "network_dir": "./network_checkpoint_potential_flow_2d",
                "decay_steps": 4000,
                "max_steps": 400000,
                "layer_size": 100,
            }
        )

if __name__ == "__main__":
    ctr = ModulusController(PotentialSolver)
    ctr.run()